#!/usr/bin/env python
#Create Hogg plots for demonstration purposes

#    Copyright 2008 Tim Weinzirl (timw@astro.as.utexas.edu)

#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

from cosmos import *
from pylab import *

rc('text',usetex=True)
rc('legend',fontsize=8)

z=arange(0.1,5,0.1)
a=Cosmo(omega0=1.0,omegalambda=0.0,verbose=0)  #solid line
b=Cosmo(omega0=0.05,omegalambda=0.0,verbose=0) #dotted line
c=Cosmo(omega0=0.2,omegalambda=0.8,verbose=0)  #dashed line

dh1,th1=a.dhubble(),a.thubble()
dh2,th2=b.dhubble(),b.thubble()
dh3,th3=c.dhubble(),c.thubble()

#legend markers
s1=r'Einstein-de Sitter ($\Omega_M$, $\Omega_\Lambda=1, 0$)'
s2=r'Low density ($\Omega_M$, $\Omega_\Lambda=0.05, 0$)'
s3=r'High $\Lambda$ ($\Omega_M$, $\Omega_\Lambda=0.2, 0.8$)'

#figure dimensions
ioff()
figure(figsize=(6,6))

#plot 1: z vs dcomovingtransverse
clf()
y1,y2,y3=[],[],[]
for zz in z:
 y1.append(a.dcomovingtransverse(zz)/dh1)
 y2.append(b.dcomovingtransverse(zz)/dh2)
 y3.append(c.dcomovingtransverse(zz)/dh3)
plot(z,y1,'k-',z,y2,'k:',z,y3,'k--')
xlim(0,5)
ylim(0,3)

gca().xaxis.set_major_locator(MultipleLocator(1))
gca().xaxis.set_minor_locator(MultipleLocator(0.2))
gca().yaxis.set_major_locator(MultipleLocator(1))
gca().yaxis.set_minor_locator(MultipleLocator(0.2))

title('Plot 1')
xlabel(r'Redshift $z$')
ylabel(r'Proper Motion Distance $D_M/D_H$')
legend((s1,s2,s3),loc='best')
savefig('pyhogg_plot1.png')

#plot 2: z vs dangular/dh
clf()
y1,y2,y3=[],[],[]
for zz in z:
 y1.append(a.dangular(zz)/dh1)
 y2.append(b.dangular(zz)/dh2)
 y3.append(c.dangular(zz)/dh3)
plot(z,y1,'k-',z,y2,'k:',z,y3,'k--')
xlim(0,5)
ylim(0,0.50)

gca().xaxis.set_major_locator(MultipleLocator(1))
gca().xaxis.set_minor_locator(MultipleLocator(0.2))
gca().yaxis.set_major_locator(MultipleLocator(0.1))
gca().yaxis.set_minor_locator(MultipleLocator(0.02))

title('Plot 2')
xlabel(r'Redshift $z$')
ylabel(r'Angular Diameter Distance $D_A/D_H$')
legend((s1,s2,s3),loc='best')
savefig('pyhogg_plot2.png')

#plot 3: z vs dluminosity/dh
clf()
y1,y2,y3=[],[],[]
for zz in z:
 y1.append(a.dluminosity(zz)/dh1)
 y2.append(b.dluminosity(zz)/dh2)
 y3.append(c.dluminosity(zz)/dh3)
plot(z,y1,'k-',z,y2,'k:',z,y3,'k--')
xlim(0,5)
ylim(0,16)

gca().xaxis.set_major_locator(MultipleLocator(1))
gca().xaxis.set_minor_locator(MultipleLocator(0.2))
gca().yaxis.set_major_locator(MultipleLocator(5))
gca().yaxis.set_minor_locator(MultipleLocator(1))

title('Plot 3')
xlabel(r'Redshift $z$')
ylabel(r'Luminosity Distance $D_L/D_H$')
legend((s1,s2,s3),loc='best')
savefig('pyhogg_plot3.png')

#plot 4: z vs dmodulus(z) +5*log10(red100)
clf()
y1,y2,y3=[],[],[]
for zz in z:
 y1.append(a.dmodulus(zz)+5*log10(a.h100))
 y2.append(b.dmodulus(zz)+5*log10(a.h100))
 y3.append(c.dmodulus(zz)+5*log10(a.h100))
plot(z,y1,'k-',z,y2,'k:',z,y3,'k--')
xlim(0,5)
ylim(40,50)

gca().xaxis.set_major_locator(MultipleLocator(1))
gca().xaxis.set_minor_locator(MultipleLocator(0.2))
gca().yaxis.set_major_locator(MultipleLocator(2))
gca().yaxis.set_minor_locator(MultipleLocator(0.5))

title('Plot 4')
xlabel(r'Redshift $z$')
ylabel(r'Distance Modulus DM + 5 log $h$ (mag)')
legend((s1,s2,s3),loc='best')
savefig('pyhogg_plot4.png')

#plot 5: z vs dvcomoving(z)/(dhubble()**3)
clf()
y1,y2,y3=[],[],[]
for zz in z:
 y1.append(a.dvcomoving(zz)/dh1**3)
 y2.append(b.dvcomoving(zz)/dh2**3)
 y3.append(c.dvcomoving(zz)/dh3**3)
plot(z,y1,'k-',z,y2,'k:',z,y3,'k--')
xlim(0,5)
ylim(0,1.1)

gca().xaxis.set_major_locator(MultipleLocator(1))
gca().xaxis.set_minor_locator(MultipleLocator(0.2))
gca().yaxis.set_major_locator(MultipleLocator(0.2))
gca().yaxis.set_minor_locator(MultipleLocator(0.05))

title('Plot 5')
xlabel(r'Redshift $z$')
ylabel(r'Comoving Volume Element $[1/D_H]^3$ $dV_c/dz/d\Omega$')
legend((s1,s2,s3),loc='best')
savefig('pyhogg_plot5.png')

#plot 6: age and lookback time
clf()
y1,y2,y3,y4,y5,y6=[],[],[],[],[],[]
for zz in z:
 cage1=a.getage(zz)/th1
 cage2=b.getage(zz)/th2
 cage3=c.getage(zz)/th3
 clook1=a.getage(0)/th1-cage1
 clook2=b.getage(0)/th2-cage2
 clook3=c.getage(0)/th2-cage3
 y1.append(cage1)
 y2.append(cage2)
 y3.append(cage3)
 y4.append(clook1)
 y5.append(clook2)
 y6.append(clook3)
plot(z,y1,'k-',z,y2,'k:',z,y3,'k--',z,y4,'k-',z,y5,'k:',z,y6,'k--')
xlim(0,5)
ylim(0,1.2)

gca().xaxis.set_major_locator(MultipleLocator(1))
gca().xaxis.set_minor_locator(MultipleLocator(0.2))
gca().yaxis.set_major_locator(MultipleLocator(0.2))
gca().yaxis.set_minor_locator(MultipleLocator(0.05))

title('Plot 6')
xlabel(r'Redshift $z$')
ylabel(r'Lookback Time $t_L/t_H$ and age $t/t_H$')
legend((s1,s2,s3),loc='best')
savefig('pyhogg_plot6.png')

#plot 7: z vs (1+z)**2/epeebles(z)
clf()
y1,y2,y3=[],[],[]
for zz in z:
 y1.append((1+zz)**2/a.epeebles(zz))
 y2.append((1+zz)**2/b.epeebles(zz))
 y3.append((1+zz)**2/c.epeebles(zz))
plot(z,y1,'k-',z,y2,'k:',z,y3,'k--')
xlim(0,5)
ylim(0,6)

gca().xaxis.set_major_locator(MultipleLocator(1))
gca().xaxis.set_minor_locator(MultipleLocator(0.2))
gca().yaxis.set_major_locator(MultipleLocator(2))
gca().yaxis.set_minor_locator(MultipleLocator(0.5))

title('Plot 7')
xlabel(r'Redshift $z$')
ylabel(r'Dimensionless Intersection Probability $dP/dz$')
legend((s1,s2,s3),loc='best')
savefig('pyhogg_plot7.png')
